import networkx
import scipy.sparse as sprs
import numpy as np
#import scipy.sparse
import scipy.sparse as sprs
import copy 
import random as r_n
import pylab
from random import randint 
from matplotlib import mlab
from scipy.stats import bernoulli
from scipy.optimize import minimize
import pandas as pd
from scipy.sparse import csr_matrix, lil_matrix, hstack, vstack
import random
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
from sklearn import preprocessing 
from sklearn.preprocessing import normalize
from random import choice
from scipy.stats import binom
import os
from sklearn.model_selection import ParameterGrid
from multiprocessing import Process
from IPython.display import clear_output
from statsmodels.stats.proportion import proportion_confint
from IPython.display import clear_output
import math
from scipy import integrate
import seaborn as sns
from IPython.display import clear_output
import time
from sympy import Symbol
from sympy import *
from scipy.optimize import minimize, rosen, rosen_der
import pickle



def add_network(graph_model, n, parameters):
    
    """
    create social network G using a predefined random graph model
    arguments:
    graph_model - the random graph model to be applied
    n - the number of vertices
    parameters - parameter(s) of the random graph model 
    Created social network must be connected - otherwise the algorithm is repeated 
    return adjacency matrix and positions of nodes (important only for random geometric networks, 
    for which positions of nodes are esstinal for the network building and vizualization)
    positions are defined using spectral layout (excepting the case of the cycle graph for which circular 
    layout is applied)
    """
    
    if graph_model == 'r_g':  #random geometric graph model
        
        start = 1
        while start:
            pos = {i: (random.uniform(0, 2), random.uniform(0, 2)) for i in range(n)}
            graph = networkx.random_geometric_graph(n, 2*parameters, pos=pos)
            if networkx.is_connected(graph):
                start = 0
        A = networkx.adjacency_matrix(graph)     
        return A, pos

    elif graph_model == 'gnm':  #Erdős–Rényi model
        
        start = 1
        while start:
            graph = networkx.gnm_random_graph(n, parameters)
            if networkx.is_connected(graph):
                start = 0
        A = networkx.adjacency_matrix(graph) 
        pos = networkx.spectral_layout(graph)

        return A, pos

    elif graph_model == 'BA':  #Barabasi-Albert model
        
        start = 1
        while start:
            graph = networkx.barabasi_albert_graph(n, parameters)
            if networkx.is_connected(graph):
                start = 0
        A = networkx.adjacency_matrix(graph) 
        pos = networkx.spectral_layout(graph)

        return A, pos

    elif graph_model == 'WS':  #Watts-Strogatz model
        
        start = 1
        while start:
            graph = networkx.watts_strogatz_graph(n, parameters[0], parameters[1])
            if networkx.is_connected(graph):
                start = 0
        A = networkx.adjacency_matrix(graph) 
        pos = networkx.spectral_layout(graph)

        return A, pos
    
    elif graph_model == 'complete':  #Complete graph
        
        A = np.ones((n, n))
        G=networkx.from_numpy_matrix(A)
        A = networkx.adjacency_matrix(G) 
        pos = networkx.spectral_layout(G)     
        
        return A, pos
    
    elif graph_model == 'cycle':  #Cycle graph
        
        graph = networkx.cycle_graph(n)
        A = networkx.adjacency_matrix(graph)
        pos = networkx.circular_layout(graph)
        
        
        return A, pos
        
    else:
        
        raise ValueError("Wrong network type")

        
def polarization_index(opinions):
    
    """
    compute dissimilarity coefficient for a given opinion distribution
    """
    
    distr = []
    
    for i in range(len(opinions)):
        for j in range(i):
            distr.append(abs(opinions[i] - opinions[j]))
            
    pol = np.array(distr).std()
    
    
    return pol


def assortativity(opinions, G):  
    
    """
    assortativity coefficient
    opinion values are assuemd to be nonnegative integers like 0, 1, 2, 3, 4, ...
    """
    
    #labels = opinions.astype(int)
    
    labels_dict = {k: opinions[k] for k in range(len(opinions))}
    networkx.set_node_attributes(G, labels_dict, 'opinion')

    return networkx.numeric_assortativity_coefficient(G, 'opinion')




def change_op(agent, transition_matrix, current_index, friend_index, possible_ops, opinion_indices, Ops_indices, Ops, public_opinion):
    
    """
    change opinion of a given agent 
    """
    
    transition_probs = transition_matrix[current_index][:, friend_index]
    
    new_op_index = np.random.choice(a = opinion_indices, p = transition_probs)
    new_op = possible_ops[new_op_index]
    
    Ops_indices[agent] = new_op_index
    Ops[agent] = new_op
    

    public_opinion[current_index][-1] -= 1 
    public_opinion[new_op_index][-1] += 1 
    
    return Ops, Ops_indices, public_opinion


def change_social_network(agent, current_index, friend, frs, agent_indices, Ops_indices, G, th_value):
    
    """
    agent tries to delete connection with friend and make new connection
    """
    
    #find agents wit sufficiently similar opinions
    
    same_ops_agnts = agent_indices[abs(Ops_indices - current_index) <= th_value]
    
    #same_ops_agnts = indices[abs(np.where(Ops[:, None] == possible_ops[None, :])[1] - current_index) <= th_value]
    
    #sourt out those who are already friends
    potential_frs = list(set(same_ops_agnts) - set(frs))
    
    if potential_frs:  #if this set is not empty - change social network
        new_fr = np.random.choice(potential_frs)
        G.remove_edge(agent, friend)
        G.add_edge(agent, new_fr)
    else:
        new_fr = -1
    
    
    return G, new_fr
    
 

def initial_opinions(N, m, in_ops_type, possible_ops):
    
    """
    initialize initial opinions (variable in_ops_type represents different strategies of opinion initialization)
    """
    
    if in_ops_type[0] == 'uniform':
        p_vector = np.ones(m) / m
        
    elif in_ops_type[0] == 'right-peak':
        p_vector = np.zeros(m)
        p_vector[-1] += 1
        
    elif in_ops_type[0] == 'center-peak':
        p_vector = np.zeros(m)
        middle_index = int(m/2)
        p_vector[middle_index] += 1
        
    elif in_ops_type[0] == 'left-peak':
        p_vector = np.zeros(m)
        p_vector[0] += 1

        
    elif in_ops_type[0] == 'manual': 
        
        p_vector = in_ops_type[1]
        #p_vector = np.array([0.07, 0.79, 0.14])  #only for m=3
        #p_vector = np.array([0.2, 0.8])  #only for m=2
        #p_vector = np.array([0.07, 0.75, 0.18])  #only for m=3
        
    else:
        
        raise ValueError("Wrong initial opinion distribution type")
        
    initial_ops = np.random.choice(a=possible_ops, p=p_vector, size=N)
    
    return initial_ops
    
    
    

def experiment_one_matrix(N, m, transition_matrix, in_ops_type, possible_ops, topologies, topology_id, selectivity_rate, pers_rate, th_value, t_freq, T):
    
    """
    simulation run for Models 1 and 2 employing one transition matrix
    arguments:
    N - number of agents
    m - number of opinion values
    transition_matrix - 
    in_ops_type - cofiguration of initial opinions
    possible_ops - opinion values
    topologies - the set of possible network topologies
    topology_id - the concrete topology to be used in experiment
    selectivity_rate - \gamma from Model 2
    pers_rate - \delta from Model 2
    th_value - \Delta from Model 2
    t_freq - metrics are calculated one time per t_freq iterations
    T - duration of experiment
    """
    
    #load transition matrix
    #transition_matrix = []
    #for i in range(m):
    #    transition_matrix.append(np.load(f'Transition matrices/Matrix_{m}_{i+1}.npy'))
        #transition_matrix.append(np.load(f'Transition matrices/Matrix_{m}_{i+1}_second_iteration.npy'))

    #number of agents 
    N = topologies[topology_id][0]  
    
    #agent indices
    agent_indices = np.arange(N)
    
    #opinion indices - indices of possible opinion values
    opinion_indices = np.arange(m)
    
    #######################
    
    #create social network
    A, pos = add_network(topologies[topology_id][2], topologies[topology_id][0], topologies[topology_id][1])
    
    #initialize initial opinions (variable in_ops_type represents different strategies of opinion initialization)
    initial_ops = initial_opinions(N, m, in_ops_type, possible_ops)
    
    #create object graph
    G = networkx.from_scipy_sparse_matrix(A)

    
    #initialize variables y_i that represent public opinion
    public_opinion = [[] for j in range(m)]
    for j in range(m): 
        public_opinion[j].append(sum(initial_ops == possible_ops[j]))
    
    #here we will store time moments at which we will make measurements
    t_measurements = []
    
    #empty lists, which we will append with new measurements
    pols = []  #polarization values
    assrts = []  #assortativity values
    components_number = []  #number of connected components 
    adj_matrices = []
    
    
    #opinion vector
    Ops = np.array(initial_ops)
    
    #indices vector
    Ops_indices = np.array([opinion_indices[possible_ops == Ops[i]].min() for i in range(N)])
    
    for t in tqdm(range(1, T+1)):
        
        for j in range(m):
            public_opinion[j].append(public_opinion[j][-1])
        
        
        if (t % t_freq == 0) | (t == 1) | (t == T-1):  #make measurements with frequency t_freq
            
            #make measurements
            t_measurements.append(t)
            pols.append(polarization_index(Ops))
            
            if topology_id == 'complete': 
                assrts.append(-3)
                components_number.append(-3)
                adj_matrices.append(A)
            else:
                assrts.append(assortativity(Ops_indices, G))
                components_number.append(networkx.number_connected_components(G))
                A = networkx.adjacency_matrix(G) 
                adj_matrices.append(A)
            
            #A = networkx.adjacency_matrix(G) 
            #adj_matrices.append(A)
            

        agent = np.random.choice(agent_indices)  #take a random agent
        agent_op = Ops[agent]  #take agent's opinion
        current_index = Ops_indices[agent] #take index of agent's opinion  
        
        frs = [l for l in G.neighbors(agent)]  #agent's friends
        
        if frs:  #if the set of friends is not empty
            
            friend = np.random.choice(frs)  #take a random friend
            friend_op = Ops[friend]  #random friend's opinion
            friend_index = Ops_indices[friend] #index of random friend's opinion

            if (abs(current_index-friend_index) > th_value):   #if opinions are too distand

                if (random.random() > pers_rate):  #with probability (1 - pers_rate) communication is allowed
                
                #to get probability=1, one should take pers_rate=-0.01
                #to get probability=0, one should take pers_rate=1.01
                
                
                    #with probability selectivity_rate the network evolves
                    if (random.random() > (1-selectivity_rate)):  
                    
                    #to get probability=0, one should take selectivity_rate=-0.01
                    #to get probability=1, one should take selectivity_rate=1.01
                        
                        G, new_fr = change_social_network(agent, 
                                                          current_index, 
                                                          friend, 
                                                          frs, 
                                                          agent_indices, 
                                                          Ops_indices, 
                                                          G, th_value)
                        
                    else: #with probability (1 - selectivity_rate) agent's opinion changes
                        
                        Ops, Ops_indices, public_opinion = change_op(agent, 
                                                                     transition_matrix, 
                                                                     current_index, 
                                                                     friend_index, 
                                                                     possible_ops, 
                                                                     opinion_indices, Ops_indices, Ops, public_opinion)
                
                #else:  #with probability pers_rate communication is prohibited
                    
            else:  #if opinion are close enough, then communication is always allowed
                
                Ops, Ops_indices, public_opinion = change_op(agent, 
                                                             transition_matrix, 
                                                             current_index, 
                                                             friend_index, 
                                                             possible_ops, 
                                                             opinion_indices, Ops_indices, Ops, public_opinion)             
        
    return pos, Ops, Ops_indices, public_opinion, pols, assrts, components_number, adj_matrices, t_measurements




def experiment_two_matrices(N, m, matrix1, matrix2, in_ops_type, possible_ops, topologies, topology_id, selectivity_rate, pers_rate, th_value, t_freq, T):
    
    """
    simulation run for Models 1 and 2 employing two transition matrix
    arguments:
    N - number of agents
    m - number of opinion values
    matrix1 - the first transitio matrix
    matrix2 - the second transition matrix
    in_ops_type - cofiguration of initial opinions
    possible_ops - opinion values
    topologies - the set of possible network topologies
    topology_id - the concrete topology to be used in experiment
    selectivity_rate - \gamma from Model 2
    pers_rate - \delta from Model 2
    th_value - \Delta from Model 2
    t_freq - metrics are calculated one time per t_freq iterations
    T - duration of experiment
    """
    
    #number of agents 
    N = topologies[topology_id][0]  
    
    #agent indices
    agent_indices = np.arange(N)
    
    #opinion indices - indices of possible opinion values
    opinion_indices = np.arange(m)
    
    #######################
    
    #create social network
    A, pos = add_network(topologies[topology_id][2], topologies[topology_id][0], topologies[topology_id][1])
    
    #initialize initial opinions (variable in_ops_type represents different strategies of opinion initialization)
    initial_ops = initial_opinions(N, m, in_ops_type, possible_ops)
    
    #create object graph
    G = networkx.from_scipy_sparse_matrix(A)

    
    #initialize variables y_i that represent public opinion
    public_opinion = [[] for j in range(m)]
    for j in range(m): 
        public_opinion[j].append(sum(initial_ops == possible_ops[j]))
    
    #here we will store time moments at which we will make measurements
    t_measurements = []
    
    #empty lists, which we will append with new measurements
    pols = []  #polarization values
    assrts = []  #assortativity values
    components_number = []  #number of connected components 
    adj_matrices = []
    
    
    #opinion vector
    Ops = np.array(initial_ops)
    
    #indices vector
    Ops_indices = np.array([opinion_indices[possible_ops == Ops[i]].min() for i in range(N)])
    
    for t in tqdm(range(1, T+1)):
        
        if t < T/2:
            transition_matrix = matrix1
        else:
            transition_matrix = matrix2
        
        for j in range(m):
            public_opinion[j].append(public_opinion[j][-1])
        
        
        if (t % t_freq == 0) | (t == 1) | (t == T-1):  #make measurements with frequency t_freq
            
            #make measurements
            t_measurements.append(t)
            pols.append(polarization_index(Ops))
            
            if topology_id == 'complete': 
                assrts.append(-3)
                components_number.append(-3)
                adj_matrices.append(A)
            else:
                assrts.append(assortativity(Ops_indices, G))
                components_number.append(networkx.number_connected_components(G))
                A = networkx.adjacency_matrix(G) 
                adj_matrices.append(A)
            
        agent = np.random.choice(agent_indices)  #take a random agent
        agent_op = Ops[agent]  #take agent's opinion
        current_index = Ops_indices[agent] #take index of agent's opinion  
        
        frs = [l for l in G.neighbors(agent)]  #agent's friends
        
        if frs:  #if the set of friends is not empty
            
            friend = np.random.choice(frs)  #take a random friend
            friend_op = Ops[friend]  #random friend's opinion
            friend_index = Ops_indices[friend] #index of random friend's opinion

            if (abs(current_index-friend_index) > th_value):   #if opinions are too distand

                if (random.random() > pers_rate):  #with probability (1 - pers_rate) communication is allowed
                
                #to get probability=1, one should take pers_rate=-0.01
                #to get probability=0, one should take pers_rate=1.01
                
                
                    #with probability selectivity_rate the network evolves
                    if (random.random() > (1-selectivity_rate)):  
                    
                    #to get probability=0, one should take selectivity_rate=-0.01
                    #to get probability=1, one should take selectivity_rate=1.01
                        
                        G, new_fr = change_social_network(agent, 
                                                          current_index, 
                                                          friend, 
                                                          frs, 
                                                          agent_indices, 
                                                          Ops_indices, 
                                                          G, th_value)
                        
                    else: #with probability (1 - selectivity_rate) agent's opinion changes
                        
                        Ops, Ops_indices, public_opinion = change_op(agent, 
                                                                     transition_matrix, 
                                                                     current_index, 
                                                                     friend_index, 
                                                                     possible_ops, 
                                                                     opinion_indices, Ops_indices, Ops, public_opinion)
                
                #else:  #with probability pers_rate communication is prohibited
                    
            else:  #if opinion are close enough, then communication is always allowed
                
                Ops, Ops_indices, public_opinion = change_op(agent, 
                                                             transition_matrix, 
                                                             current_index, 
                                                             friend_index, 
                                                             possible_ops, 
                                                             opinion_indices, Ops_indices, Ops, public_opinion)             
        
    return pos, Ops, Ops_indices, public_opinion, pols, assrts, components_number, adj_matrices, t_measurements


def save_obj(obj, path_to_save, name):
    with open(f'{path_to_save}/{name}.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(path_to_load, name):
    with open(f'{path_to_load}/{name}.pkl', 'rb') as f:
        return pickle.load(f)


def save(main_folder, arguments, results):
    
    """
    save experiment 
    """
    
    N, m, in_ops_type, topology_id, selectivity_rate, pers_rate, th_value, t_freq, T = arguments
    
    #pos, Ops, public_opinion, pols, assrts, components_number, adj_matrices, t_measurements = results
    
    pos, Ops, Ops_indices, public_opinion, pols, assrts, components_number, adj_matrices, t_measurements = results
    

    #prepare space for storing experimental data
    
    #main_folder = 'Experiments_Model_1'
    
    files = os.listdir()
    if main_folder not in files:
        os.mkdir(main_folder)
        
    files = os.listdir(main_folder)
    
    if len(in_ops_type) == 1:
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type}'
    else:
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type[0]}'
        for i in range(len(in_ops_type[1])):
            folder += f'_{in_ops_type[1][i]}'
    
    folder += f'_topology_id={topology_id}'
    folder += f'_selectivity_rate={selectivity_rate}_pers_rate={pers_rate}_th_value={th_value}'
    folder += f'_t_freq={t_freq}_T={T}'

    if folder not in files:

        os.mkdir(f'{main_folder}/{folder}')


    files = os.listdir(f'{main_folder}/{folder}')

    #identify the counting number of the experiment
    counter = 1

    for file in files:
        if len(file) > 3:
            if file[:4] == 'exp_':
                counter += 1

    path = f'{main_folder}/{folder}/exp_{counter}'

    os.mkdir(path)

    #save files
    save_obj(pos, path, 'pos')
    
    #A_0 = networkx.adjacency_matrix(G_0)
    #A = networkx.adjacency_matrix(G)
    
    sprs.save_npz(f'{path}/A_initial.npz', adj_matrices[0])
    sprs.save_npz(f'{path}/A_final.npz', adj_matrices[-1])
    
    np.save(f'{path}/Ops.npy', Ops)
    
    for i in range(len(public_opinion)):
        np.save(f'{path}/Y_{i+1}.npy', public_opinion[i])
    
    np.save(f'{path}/pols.npy', pols)
    np.save(f'{path}/assrts.npy', assrts)
    np.save(f'{path}/components_number.npy', components_number)
    np.save(f'{path}/t_measurements.npy', t_measurements)
    #Posts_base.to_csv(f'{path}/Posts_base.csv', index=False)
    #Actions_base.to_csv(f'{path}/Actions_base.csv', index=False)
    
    return 0



def load(main_folder, N, m, in_ops_type, topology_id, selectivity_rate, pers_rate, th_value, t_freq, T, counter):
    
    
    """
    load experiment
    """
    
    if len(in_ops_type) == 1:
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type}'
    else:
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type[0]}'
        for i in range(len(in_ops_type[1])):
            folder += f'_{in_ops_type[1][i]}'
    
    folder += f'_topology_id={topology_id}'
    folder += f'_selectivity_rate={selectivity_rate}_pers_rate={pers_rate}_th_value={th_value}'
    folder += f'_t_freq={t_freq}_T={T}'

    path = f'{main_folder}/{folder}/exp_{counter}'
    
    pos = load_obj(path, 'pos')
    A_0 = sprs.load_npz(f'{path}/A_initial.npz')
    A = sprs.load_npz(f'{path}/A_final.npz')
    
    Ops = np.load(f'{path}/Ops.npy')
    
    public_opinion = []
    for i in range(m):
        public_opinion.append(np.load(f'{path}/Y_{i+1}.npy'))  
    
    pols = np.load(f'{path}/pols.npy')
    assrts = np.load(f'{path}/assrts.npy')
    components_number = np.load(f'{path}/components_number.npy')
    t_measurements = np.load(f'{path}/t_measurements.npy')

    return pos, Ops, public_opinion, pols, assrts, components_number, A_0, A, t_measurements 

def save_noise(main_folder, arguments, results, alpha, beta, phi, psi):
    
    """
    save experiment (noise in transition matrix)
    """
    
    N, m, in_ops_type, topology_id, selectivity_rate, pers_rate, th_value, t_freq, T = arguments
    
    #pos, Ops, public_opinion, pols, assrts, components_number, adj_matrices, t_measurements = results
    
    pos, Ops, Ops_indices, public_opinion, pols, assrts, components_number, adj_matrices, t_measurements = results
    

    #prepare space for storing experimental data
    
    #main_folder = 'Experiments_Model_1'
    
    files = os.listdir()
    if main_folder not in files:
        os.mkdir(main_folder)
        
    files = os.listdir(main_folder)
    
    if len(in_ops_type) == 1:
        
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type}'
    else:
        
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type[0]}'
        
        for i in range(len(in_ops_type[1])):
            
            folder += f'_{in_ops_type[1][i]}'
    
    folder += f'_topology_id={topology_id}'
    folder += f'_selectivity_rate={selectivity_rate}_pers_rate={pers_rate}_th_value={th_value}'
    folder += f'_t_freq={t_freq}_T={T}'
    folder += f'_alpha={alpha}_beta={beta}_phi={phi}_psi={psi}'

    if folder not in files:

        os.mkdir(f'{main_folder}/{folder}')


    files = os.listdir(f'{main_folder}/{folder}')

    #identify the counting number of the experiment
    counter = 1

    for file in files:
        if len(file) > 3:
            if file[:4] == 'exp_':
                counter += 1

    path = f'{main_folder}/{folder}/exp_{counter}'

    os.mkdir(path)

    #save files
    save_obj(pos, path, 'pos')
    
    #A_0 = networkx.adjacency_matrix(G_0)
    #A = networkx.adjacency_matrix(G)
    
    sprs.save_npz(f'{path}/A_initial.npz', adj_matrices[0])
    sprs.save_npz(f'{path}/A_final.npz', adj_matrices[-1])
    
    np.save(f'{path}/Ops.npy', Ops)
    
    for i in range(len(public_opinion)):
        np.save(f'{path}/Y_{i+1}.npy', public_opinion[i])
    
    np.save(f'{path}/pols.npy', pols)
    np.save(f'{path}/assrts.npy', assrts)
    np.save(f'{path}/components_number.npy', components_number)
    np.save(f'{path}/t_measurements.npy', t_measurements)
    #Posts_base.to_csv(f'{path}/Posts_base.csv', index=False)
    #Actions_base.to_csv(f'{path}/Actions_base.csv', index=False)
    
    return 0


def load_noise(main_folder, N, m, in_ops_type, topology_id, selectivity_rate, pers_rate, th_value, t_freq, T, alpha, beta, phi, psi, counter):

    
    """
    load experiment (noise in transition matrix)
    """
        #prepare space for storing experimental data
    
    #main_folder = 'Experiments_Model_1'
    
    if len(in_ops_type) == 1:
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type}'
    else:
        folder = f'N={N}_m={m}_in_ops_type={in_ops_type[0]}'
        for i in range(len(in_ops_type[1])):
            folder += f'_{in_ops_type[1][i]}'
    
    folder += f'_topology_id={topology_id}'
    folder += f'_selectivity_rate={selectivity_rate}_pers_rate={pers_rate}_th_value={th_value}'
    folder += f'_t_freq={t_freq}_T={T}'
    folder += f'_alpha={alpha}_beta={beta}_phi={phi}_psi={psi}'

    path = f'{main_folder}/{folder}/exp_{counter}'
    
    pos = load_obj(path, 'pos')
    A_0 = sprs.load_npz(f'{path}/A_initial.npz')
    A = sprs.load_npz(f'{path}/A_final.npz')
    
    Ops = np.load(f'{path}/Ops.npy')
    
    public_opinion = []
    for i in range(m):
        public_opinion.append(np.load(f'{path}/Y_{i+1}.npy'))  
    
    pols = np.load(f'{path}/pols.npy')
    assrts = np.load(f'{path}/assrts.npy')
    components_number = np.load(f'{path}/components_number.npy')
    t_measurements = np.load(f'{path}/t_measurements.npy')

    return pos, Ops, public_opinion, pols, assrts, components_number, A_0, A, t_measurements

###Modules for perfoming symbolic computations

def right_part(x, f, m, P):  #function g from system (1) in the main manuscript
    
    """
    x should initially have m-1 components
    """
    
    x = list(x)
    x = x + [1-sum(x)]
    x = tuple(x)

    res = 0
    
    for s in range(m):
        for l in range(m):
            for k in range(m):
                
                res = res + x[s]*x[l]*P[s][l][k]*(int((k == f) - int(s == f)))
    
    return res

def all_right_parts(x):  #objective function from optimization problem (6) in the main manuscript
    
    res = 0
    
    for f in range(m-1):
        
        res = res + right_part(x, f, m, P) ** 2
        
    return res